"""
进行deepsort追踪
"""
import argparse
import csv
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # 解决多线程引起的错误
from utils.general import xyxy2xywh
from deep_sort_pytorch.utils.parser import get_config # 读取deepsort配置
from deep_sort_pytorch.deep_sort import DeepSort

import cv2
import torch
import numpy as np
import pickle
from PIL import Image
torch.set_num_threads(1)

# 存储每帧的检测和跟踪结果
dicts = []

# 执行yolov5检测和deepsort跟踪流程
def detect(opt):
  source = opt.source
  stride = 32
  pt = True
  jit = False

  # 读取deepsort配置文件，初始化追踪参数
  cfg = get_config()
  cfg.merge_from_file(opt.config_deepsort)

  # 从pkl中加载检测坐标及置信度信息
  with open(r'D:\file\postgrad\experiment\bird_ava_dataset\avaMin_dense_proposals_train.pkl', 'rb') as f:
    info = pickle.load(f, encoding='iso-8859-1')

  # 记录当前处理的文件名称
  tempFileName = ''

  # 定义类别区间与类别名称的映射（如果没分类别，这部分可以不加）
  def get_category_by_filename(filename):
    temp_id = int(filename)
    if 1 <= temp_id <= 247:
      return 0  # silver pheasant
    elif 248 <= temp_id <= 330:
      return 1  # tragopan
    if 331 <= temp_id <= 698:
      return 0  # silver pheasant
    else:
      return -1  # 如果没有匹配的ID，返回 -1 作为异常
  # 循环 pkl中的信息
  for i in info:
    dets = info[i]  # 存储检测坐标和置信度
    tempName = i.split(',') # 提取当前文件和帧数信息

    # 获取类别ID
    category = get_category_by_filename(tempName[0])
    # 如果没有匹配的类别，跳过此帧
    if category == -1:
      print(f"Warning: No category found for filename {tempName[0]}. Skipping this file.")
      continue

    # 创建新的deepsort对象，对新文件进行追踪
    if tempName[0] != tempFileName:
      deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
                          max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
                          max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
                          max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
                          use_cuda=True)
      tempFileName = tempName[0] # 更新当前文件名

    # 读取当前帧图像并计算坐标
    im0Path = os.path.join(
      source, tempName[0], f"{tempName[0]}_{str(int(tempName[1])).zfill(6)}.jpg"
    )
    print(f"Reading image from: {im0Path}")
    tempImg = cv2.imread(im0Path)
    # 检查图像是否加载成功
    if tempImg is None:
      print(f"Error: Unable to load image at {im0Path}. Skipping this frame.")
      continue  # 跳过当前帧
    # 获取图片的大小
    imgsz = tempImg.shape

    # 将pkl中的检测坐标从xyxy转换为xywh
    xyxys = torch.FloatTensor(len(dets), 4)
    confs = torch.FloatTensor(len(dets))
    clss = torch.FloatTensor(len(dets)) # 存储类别
    for index, det in enumerate(dets):
      xyxys[index][0] = det[0] * imgsz[1]
      xyxys[index][1] = det[1] * imgsz[0]
      xyxys[index][2] = det[2] * imgsz[1]
      xyxys[index][3] = det[3] * imgsz[0]
      # confs[index] = (float(det[4]))
      confs[index] = 0.98
      clss[index] = category

    xywhs = xyxy2xywh(xyxys) # 坐标格式转换

    # 读取图像并更新deepsort跟踪器
    im0Path = os.path.join(source, tempName[0], f"{tempName[0]}_{str(int(tempName[1])).zfill(6)}.jpg")
    im0 = np.array(Image.open(im0Path))
    # 检查图像是否加载成功
    if im0 is None:
      print(f"Error: Unable to load image at {im0Path}. Skipping this frame.")
      continue  # 跳过当前帧
    outputs = deepsort.update(xywhs, confs, clss, im0) # 更新deepsort跟踪结果


    if len(outputs) > 0:
      for output in outputs:
        x1 = output[0] / imgsz[1]
        y1 = output[1] / imgsz[0]
        x2 = output[2] / imgsz[1]
        y2 = output[3] / imgsz[0]
        dict = [tempName[0], tempName[1], x1, y1, x2, y2, output[4]]
        dicts.append(dict)
    # 将所有帧追踪结果写入CSV文件
    with open(r'D:\file\postgrad\experiment\bird_ava_dataset\train_personID.csv', "w", newline='') as csvfile:
      writer = csv.writer(csvfile)
      writer.writerows(dicts)


if __name__ == '__main__':
  parser = argparse.ArgumentParser()

  parser.add_argument('--deep_sort_weights', type=str, default=r'D:\file\postgrad\experiment\code\datadeal\deep_sort_pytorch\deep_sort\deep\checkpoint\ckpt.t7',
                      help='ckpt.t7 path')
  # file/folder, 0 for webcam
  parser.add_argument('--source', type=str, default=r'D:\file\postgrad\experiment\bird_ava_dataset\videos_labelframes', help='source')
  parser.add_argument('--save-txt', action='store_true', help='save MOT compliant results to *.txt')
  # class 0 is person, 1 is bycicle, 2 is car... 79 is oven
  parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 16 17')
  parser.add_argument("--config_deepsort", type=str, default=r"D:\file\postgrad\experiment\code\datadeal\deep_sort_pytorch\configs\deep_sort.yaml")

  opt = parser.parse_args()
  with torch.no_grad():
    detect(opt)